package org.example;

import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.SystemPromptTemplate;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.io.Resource;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * AI服务类 - 处理与AI模型的交互，生成聊天响应
 * 负责构建提示、管理会话记忆和处理流式响应
 */
@Service
public class AIService {

    /**
     * ChatClient - Spring AI提供的客户端，用于与AI模型进行通信
     */
    @Autowired
    private ChatClient chatClient;

    /**
     * ChatMemory - 聊天记忆组件，用于存储和管理对话历史
     */
    @Autowired
    private ChatMemory chatMemory;

    @Autowired
    ReportResultService reportResultService;

    @Autowired
    VectorStore vectorStore;

    /**
     * 系统提示模板资源 - 从类路径加载系统提示模板文件
     * 模板定义了AI助手的行为和角色
     */
    @Value("classpath:/prompts/system-message.st")
    private Resource systemResource;

    @Value("classpath:/prompts/system-report-message.st")
    private Resource reportSystemMessageResource;

    @Value("classpath:/prompts/gen-report.st")
    private Resource genReportResource;

    /**
     * 生成AI响应的核心方法
     *
     * @param input          用户输入的提示文本
     * @param conversationId 对话会话ID，用于关联对话历史
     * @return Flux<ChatResponse> - 反应式流，包含AI模型的流式响应
     */
    public Flux<ChatResponse> generateResponse(String input, String conversationId) {

        // 创建系统提示模板实例
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource);

        // 将用户输入转换为用户消息对象
        Message userMessage = new UserMessage(input);

        // 从模板创建系统消息对象
        Message systemMessage = systemPromptTemplate.createMessage(Map.of("round", 6));

        // 构建提示对象，包含用户消息和系统消息
        Prompt prompt = new Prompt(List.of(userMessage, systemMessage));

        // 调用AI模型生成响应
        // 使用MessageChatMemoryAdvisor添加对话历史上下文
        // 使用stream()方法启用流式响应
        return chatClient
                .prompt(prompt)
                .advisors(
                        MessageChatMemoryAdvisor.builder(chatMemory)
                                .conversationId(conversationId) // 设置对话ID以关联历史消息
                                .build())
                .tools(reportResultService)
                .toolContext(Map.of("conversationId", conversationId))
                .stream() // 启用流式响应模式
                .chatResponse(); // 返回ChatResponse类型的响应流
    }

    public String generateReportResponse(String promptStr) {
        SystemPromptTemplate reportSystemPrompt =
                new SystemPromptTemplate(genReportResource);

        // 从模板创建系统消息对象
        Message systemMessage = reportSystemPrompt.createMessage(Map.of("chat_history", promptStr));

        // 构建提示对象，包含用户消息和系统消息
        Prompt prompt = new Prompt(List.of(systemMessage));

        return chatClient
                .prompt(prompt)
                .call()
                .content();
    }

    public String generateResponseBatch(String input, String conversationId) {
        // 创建系统提示模板实例
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(systemResource);

        // 将用户输入转换为用户消息对象
        Message userMessage = new UserMessage(input);

        // 从模板创建系统消息对象
        Message systemMessage = systemPromptTemplate.createMessage(Map.of("round", 4));

        // 构建提示对象，包含用户消息和系统消息
        Prompt prompt = new Prompt(List.of(userMessage, systemMessage));

        // 调用AI模型生成响应
        // 使用MessageChatMemoryAdvisor添加对话历史上下文
        // call()方法启用batch 响应
        return chatClient
                .prompt(prompt)
                .advisors(
                        MessageChatMemoryAdvisor.builder(chatMemory)
                                .conversationId(conversationId) // 设置对话ID以关联历史消息
                                .build())
                .tools(reportResultService)
                .call()
                .content();
    }

    public void addDocument(String content) {
        String keywords = chatClient.prompt("为下面的文本提取5个关键词，关键词之间以逗号隔开: " + content).call().content();
        System.out.println("Extracted keywords: " + keywords);

        Map<String, Object> metadata = new HashMap<>();
        if (keywords != null) {
            metadata = Map.of("keywords", keywords);
        }
        Document document = new Document(content, metadata);
        List<Document> documents = List.of(document);
        vectorStore.add(documents);
    }

    public List<Document> searchDocument(String query) {
        SearchRequest searchRequest = SearchRequest.builder()
                .query(query).topK(5).build();
        List<Document> results = vectorStore.similaritySearch(searchRequest);
        for (Document result : results) {
            String snappit = "";
            if (result.getText() != null) {
                snappit = result.getText().length() > 20 ? result.getText().substring(0, 20) : result.getText();
            }
            System.out.println("Found document: " + snappit + "\n" +
                    "score:" + result.getScore());
        }
        return results;
    }
}
